import os
import sys
import pandas as pd
# #受体名字
#  = sys.argv[1]
# #配体名字
#  = sys.argv[2]
# #是否block抗体（配体是抗体的话）
#  = sys.argv[3]

def auto_zdock(receptor,ligand,antibody_block):
	dir = receptor.split(".pdb")[0] + "_" + ligand.split(".pdb")[0]
	os.system("./mark_sur {} receptor.pdb".format(receptor))
	os.system("./mark_sur {} ligand.pdb".format(ligand))

	if (antibody_block == "y") | (antibody_block == "Y"):
		name_out = "receptor_block.pdb"
		list_hv = []
		list_lv = []
		# block FR 区域
		for i in range(98):
			if (i + 1 < 31-3) | ((i + 1 > 35-3) & (i + 1 < 50+3)) | ((i + 1 > 66-3) & (i + 1 < 99+3)):
				list_hv.append(str(i + 1))
		for i in range(98):
			if (i + 1 < 24-3) | ((i + 1 > 34-3) & (i + 1 < 50+3)) | ((i + 1 > 56-3) & (i + 1 < 89+3)):
				list_lv.append(str(i + 1))
		# 读取 抗体pdb文件

		pdb = open("receptor.pdb", "r")
		pdb_line = [i for i in pdb]
		pdb_line_atom = []
		for i in pdb_line:
			if i[0:4] == 'ATOM':
				pdb_line_atom.append(i)
		list_hv.append(pdb_line_atom[0][20:23].strip())
		list_lv.append(pdb_line_atom[-1][20:23].strip())
		for i in range(len(pdb_line_atom)):
			if (pdb_line_atom[i][20:23].strip() in list_hv) & (pdb_line_atom[i][23:26].strip() in list_hv):
				pdb_line_atom[i] = pdb_line_atom[i][0:55] + "19 " + pdb_line_atom[i][58:-1] + pdb_line_atom[i][-1]
			elif (pdb_line_atom[i][20:23].strip() in list_lv) & (pdb_line_atom[i][23:26].strip() in list_lv):
				pdb_line_atom[i] = pdb_line_atom[i][0:55] + "19 " + pdb_line_atom[i][58:-1] + pdb_line_atom[i][-1]
		# print(pdb_line_atom[i][54:58].strip())
		# print(pdb_line_atom[i][20:23].strip())
		pdb_w = open(name_out, "w")
		pdb_w.writelines(pdb_line_atom)
		pdb_w.close()
		pdb.close()
		print("受体抗体蛋白 已封闭")
		print("zdock 开始对接")
		os.system("./zdock -R receptor_block.pdb -L ligand.pdb -o zdock.out -D 6 -N 2000".format(dir))
	else:
		print("zdock 开始对接")
		os.system("./zdock -R receptor.pdb -L ligand.pdb -o zdock.out -D 6 -N 2000".format(dir))
	os.system("./create.pl ./zdock.out".format(dir))

	# 取 打分前100的对接复合物
	# import os
	os.system("mkdir {}".format(dir))
	for i in range(100):
		os.system("mv ./complex.{n}.pdb ./complex_{n}.pdb".format(n=i + 1, dir=dir))
	os.system("rm ./complex.*.pdb")

	os.system("chmod 777 ./*")
	os.system("chmod 777 ../*")
	list = os.listdir("./")
	list2 = []
	dic = {}
	for i in list:
		if (i[-4:-1] + i[-1] == ".pdb") & (i[0:8] == "complex_"):
			list2.append(str(i).split(".pdb")[0])
	print(list2)
	for n in list2:
		os.system(
			"pdb2pqr30 --ff AMBER --titration-state-method propka --with-ph 7 --keep-chain {n}.pdb complex_{n}.pqr --pdb-output {n}_add_H.pdb".format(
				dir=dir, n=n))
	# os.system("InterfaceAnalyzer.linuxgccrelease -s {}.pdb -fixedchains A B @pack_input_options.txt".format(n))

	os.system("ls *_add_H.pdb > list.txt")
	os.system("./zrank list.txt")

	# 打分，输出结果
	os.system("mv ./list.txt.zr.out ./zrank.txt")
	zrank = pd.read_table('./zrank.txt', sep='\t', header=None)
	zrank_add_zrank_score = pd.DataFrame({"rank-score": [i for i in range(100)]})
	zrank_ranked_combine = pd.concat([zrank, zrank_add_zrank_score], axis=1)
	zrank_ranked_combine.sort_values(by=[1], ascending=True, inplace=True)
	zrank_ranked_combine.reset_index(drop=True, inplace=True)
	#
	for i in range(100):
		zrank_ranked_combine.loc[i, "rank-score"] = int(zrank_ranked_combine.loc[i, 0].split("_")[1]) + i + 1

	zrank_ranked_combine.sort_values(by=["rank-score"], ascending=True, inplace=True)
	# print(zrank_ranked_combine)
	zrank_ranked_combine.to_csv(
		# self,
		"./zrank_{}_ranked.csv".format(dir),
		encoding="utf8",
		index=False,
		na_rep=""
	)

	#输出那个最好的那个模型的名字
	best_model_name = zrank_ranked_combine.iloc[0, 0]
	os.system("cp ./{best_model_name} ./output_result/{dir}-{best_model_name}".format(best_model_name=best_model_name,dir=dir))

#最后的工作
	os.system("mv ./complex* ./{}/".format(dir))
	os.system("mv ./*.out ./{}/".format(dir))
	os.system("mv ./zrank_{dir}_ranked.csv ./{dir}/".format(dir = dir))
	os.system("mv ./*.txt ./{}/".format(dir))
	os.system("mv ./{} ./output".format(dir))
	os.system("rm ./*.pdb")

